import subprocess
import os
import sys
import time
import threading
import shutil
import psutil

# 项目根目录、流目录和临时文件定义
PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
STREAM_DIR = os.path.join(PROJECT_DIR, 'stream')
TEMP_MKV = os.path.join(PROJECT_DIR, 'temp.mkv')  # 临时mkv文件路径，用于scrcpy录制并由ffmpeg读取

# 检查命令是否可用
def check_command_available(command):
    try:
        subprocess.run(
            [command, '--version'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        return True
    except FileNotFoundError:
        return False

# 检查ADB是否可用
def check_adb_available():
    try:
        subprocess.run(
            ['adb', 'version'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        return True
    except FileNotFoundError:
        return False

# 创建流目录
def create_stream_dir():
    if not os.path.exists(STREAM_DIR):
        os.makedirs(STREAM_DIR)
        print(f"已创建流目录: {STREAM_DIR}")
    else:
        # 按照用户要求，不清理目录中的旧文件
        print(f"流目录已存在: {STREAM_DIR}，已按照用户要求保留其中的文件")
        
        # 检查是否已有stream.m3u8文件
        playlist_path = os.path.join(STREAM_DIR, 'stream.m3u8')
        if os.path.exists(playlist_path):
            print(f"检测到已有播放列表文件: {playlist_path}")
        else:
            print(f"流目录中尚未创建播放列表文件")

# 启动scrcpy录制
def start_scrcpy_and_ffmpeg():
    try:
        # 停止可能正在运行的scrcpy和ffmpeg进程
        subprocess.run(['taskkill', '/F', '/IM', 'scrcpy.exe'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        subprocess.run(['taskkill', '/F', '/IM', 'ffmpeg.exe'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except Exception:
        pass
    
    # 删除可能存在的旧临时文件和流文件
    try:
        if os.path.exists(TEMP_MKV):
            os.unlink(TEMP_MKV)
    except Exception:
        pass
    
    # 创建流目录（确保存在）
    if not os.path.exists(STREAM_DIR):
        os.makedirs(STREAM_DIR)
        print(f"已创建流目录: {STREAM_DIR}")
    else:
        print(f"流目录已存在: {STREAM_DIR}")
    
    print("启动scrcpy录制到临时文件，然后使用ffmpeg读取该文件...")
    
    # 首先删除可能存在的临时文件
    if os.path.exists(TEMP_MKV):
        try:
            os.unlink(TEMP_MKV)
            print(f"已删除旧的临时文件: {TEMP_MKV}")
        except Exception as e:
            print(f"删除旧临时文件失败: {e}")
    
    # 配置scrcpy命令，输出到临时文件
    scrcpy_cmd = [
        'scrcpy',
        '--record', TEMP_MKV,  # 输出到临时文件
        '--record-format', 'mkv',  # 指定输出格式为mkv
        '-m', '1080',  # 最大宽度1080p
        '--no-control'  # 禁用控制功能，专注于录制
    ]
    print(f"scrcpy命令: {' '.join(scrcpy_cmd)}")
    
    try:
        # 先启动scrcpy进程，录制到临时文件
        scrcpy_process = subprocess.Popen(
            scrcpy_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        
        # 检查scrcpy进程是否真的在运行
        try:
            p = psutil.Process(scrcpy_process.pid)
            print(f"scrcpy进程已启动 (PID: {scrcpy_process.pid})，状态: {p.status()}")
        except psutil.NoSuchProcess:
            print("scrcpy进程启动后立即退出")
        
        # 等待临时文件创建并写入足够数据（增加等待时间到10秒）
        wait_time = 0
        while not os.path.exists(TEMP_MKV) or (os.path.exists(TEMP_MKV) and os.path.getsize(TEMP_MKV) < 1024 * 1024):  # 等待文件大小至少达到1MB
            if wait_time >= 10:
                break
            print(f"等待临时文件创建并写入足够数据，已等待 {wait_time} 秒...")
            if os.path.exists(TEMP_MKV):
                print(f"  当前文件大小: {os.path.getsize(TEMP_MKV)} 字节")
            time.sleep(1)
            wait_time += 1
        
        if not os.path.exists(TEMP_MKV):
            print(f"错误: 临时文件 {TEMP_MKV} 未能在10秒内创建")
            return None, None
        
        if os.path.getsize(TEMP_MKV) < 1024 * 1024:
            print(f"警告: 临时文件 {TEMP_MKV} 大小不足1MB ({os.path.getsize(TEMP_MKV)} 字节)，可能数据写入不完整")
        
        print(f"临时文件已创建且包含足够数据: {TEMP_MKV}, 大小: {os.path.getsize(TEMP_MKV)} 字节")
        
        # 为了避免读取正在写入的文件，我们创建一个临时副本供ffmpeg读取
        temp_copy = TEMP_MKV + '.copy'
        try:
            shutil.copy2(TEMP_MKV, temp_copy)
            print(f"已创建临时文件副本: {temp_copy}")
        except Exception as e:
            print(f"创建临时文件副本失败: {e}")
            # 如果复制失败，我们仍然尝试使用原始文件
            temp_copy = TEMP_MKV
        
        # 检查流目录是否存在，不存在则创建
        if not os.path.exists(STREAM_DIR):
            os.makedirs(STREAM_DIR)
            print(f"已创建流目录: {STREAM_DIR}")
        else:
            print(f"流目录已存在: {STREAM_DIR}")
            # 添加权限检查
            if os.access(STREAM_DIR, os.W_OK):
                print("流目录有写权限")
            else:
                print("警告: 流目录可能没有写权限")
        
        # 配置ffmpeg命令，从临时文件副本读取数据
        # 增加-re参数使ffmpeg以实时速度读取文件
        # 增加-fflags +genpts参数生成新的时间戳
        # 添加-hls_flags delete_segments+append_list以确保播放列表正确更新
        # 添加-vsync 2避免帧率问题
        # 添加-analyzeduration和-probesize参数解决格式检测问题
        ffmpeg_cmd = [
            'ffmpeg',
            '-analyzeduration', '100M',  # 增加分析时间
            '-probesize', '100M',        # 增加探测大小
            '-i', temp_copy,             # 从临时文件副本读取
            '-c:v', 'copy',              # 视频流直接复制，不重新编码
            '-c:a', 'aac',               # 音频转码为AAC
            '-vsync', '2',               # 避免帧率问题
            '-fflags', '+genpts',        # 生成新的时间戳
            '-re',                       # 以实时速度读取文件
            '-f', 'hls',                 # 输出格式为HLS
            '-hls_time', '2',            # 每个分片时长2秒
            '-hls_list_size', '5',       # 播放列表保留5个分片
            '-hls_playlist_type', 'event',  # 使用event类型而不是vod，确保流可以持续更新
            '-hls_flags', 'delete_segments+append_list',  # 删除旧分片并追加到播放列表
            '-y',                        # 覆盖已存在的文件
            os.path.join(STREAM_DIR, 'stream.m3u8')
        ]
        print(f"ffmpeg命令: {' '.join(ffmpeg_cmd)}")
        
        # 启动ffmpeg进程
        ffmpeg_process = subprocess.Popen(
            ffmpeg_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        
        # 检查ffmpeg进程是否真的在运行
        try:
            p = psutil.Process(ffmpeg_process.pid)
            print(f"ffmpeg进程已启动 (PID: {ffmpeg_process.pid})，状态: {p.status()}")
        except psutil.NoSuchProcess:
            print("ffmpeg进程启动后立即退出")
        
        # 检查scrcpy进程是否真的在运行
        try:
            p = psutil.Process(scrcpy_process.pid)
            print(f"scrcpy进程已启动 (PID: {scrcpy_process.pid})，状态: {p.status()}")
        except psutil.NoSuchProcess:
            print("scrcpy进程启动后立即退出")
        
        # 创建一个线程来监控scrcpy的输出，但不将正常信息视为错误
        def monitor_scrcpy():
            
            while scrcpy_process.poll() is None:
                line = scrcpy_process.stderr.readline().decode('utf-8', errors='ignore')
                if line:
                    # 只记录信息，不标记为错误
                    print(f"scrcpy输出: {line.strip()}")
            
            # 进程退出时，检查是否是异常退出
            if scrcpy_process.returncode != 0:
                print(f"scrcpy进程异常退出，返回代码: {scrcpy_process.returncode}")
                cleanup()
        
        # 启动监控线程
        scrcpy_monitor_thread = threading.Thread(target=monitor_scrcpy)
        scrcpy_monitor_thread.daemon = True
        scrcpy_monitor_thread.start()
        
        # 检查ffmpeg是否成功启动
        time.sleep(2)
        if ffmpeg_process.poll() is not None:
            stderr = ffmpeg_process.stderr.read().decode('utf-8', errors='ignore')
            print(f"错误: ffmpeg启动失败。错误输出: {stderr}")
            return None, None
        
        # 检查scrcpy是否成功启动
        if scrcpy_process.poll() is not None:
            stderr = scrcpy_process.stderr.read().decode('utf-8', errors='ignore')
            print(f"错误: scrcpy启动失败。错误输出: {stderr}")
            return None, None
        
        print("scrcpy和ffmpeg已成功启动，正在通过临时文件处理数据...")
        return scrcpy_process, ffmpeg_process
    except Exception as e:
        print(f"启动scrcpy和ffmpeg时出错: {e}")
        return None, None

# 启动WebSocket服务器
def start_websocket_server():
    print("启动WebSocket服务器...")
    ws_cmd = [
        sys.executable,
        os.path.join(PROJECT_DIR, 'websocket_server.py')
    ]
    
    return subprocess.Popen(
        ws_cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

# 启动HTTP服务器提供HTML页面
def start_http_server():
    print("启动HTTP服务器...")
    http_cmd = [
        sys.executable,
        '-m', 'http.server',
        '8000',
        '--directory', PROJECT_DIR
    ]
    
    return subprocess.Popen(
        http_cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

# 监控进程状态
def monitor_processes(processes):
    while True:
        all_processes_exited = True
        
        for name, process in processes.items():
            if process.poll() is not None:
                # 检查是否是正常退出（返回代码为0）
                if process.returncode == 0:
                    print(f"{name}进程正常退出")
                else:
                    print(f"警告: {name}进程异常终止，返回代码: {process.returncode}")
                    # 读取错误输出但不立即清理资源
                    stderr = process.stderr.read().decode('utf-8', errors='ignore')
                    print(f"{name}输出: {stderr}")
            else:
                all_processes_exited = False
        
        # 只有当所有进程都已退出时才返回
        if all_processes_exited:
            print("所有进程已退出，保持流文件不删除")
            return False
        
        # 检查流文件是否生成
        m3u8_path = os.path.join(STREAM_DIR, 'stream.m3u8')
        if os.path.exists(m3u8_path):
            print(f"流文件已生成: {m3u8_path}, 大小: {os.path.getsize(m3u8_path)} 字节")
            # 检查是否有ts分片文件
            ts_files = [f for f in os.listdir(STREAM_DIR) if f.endswith('.ts')]
            if ts_files:
                print(f"找到 {len(ts_files)} 个TS分片文件")
                # 添加一些调试信息，显示分片文件大小
                for ts_file in ts_files[:3]:  # 只显示前3个文件
                    ts_path = os.path.join(STREAM_DIR, ts_file)
                    print(f"  - {ts_file}: {os.path.getsize(ts_path)} 字节")
            else:
                print("流目录中未找到TS分片文件")
        else:
            print(f"流文件尚未生成: {m3u8_path}")
            # 添加调试信息，检查ffmpeg进程是否还在运行
            if 'ffmpeg' in processes and processes['ffmpeg'].poll() is None:
                print("ffmpeg进程仍在运行中")
            else:
                print("ffmpeg进程未在运行")
        
        time.sleep(2)

# 清理资源
def cleanup():
    print("正在清理资源...")
    
    try:
        # 停止可能正在运行的scrcpy进程
        subprocess.run(['taskkill', '/F', '/IM', 'scrcpy.exe'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except Exception:
        pass
    
    try:
        # 停止可能正在运行的ffmpeg进程
        subprocess.run(['taskkill', '/F', '/IM', 'ffmpeg.exe'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except Exception:
        pass
    
    # 尝试清理可能存在的临时文件（当前使用临时文件方式）
    if os.path.exists(TEMP_MKV):
        for i in range(3):
            try:
                os.unlink(TEMP_MKV)
                print(f"已删除临时录制文件: {TEMP_MKV}")
                break
            except Exception as e:
                print(f"清理临时文件时出错: {e}")
                if i < 2:  # 不是最后一次尝试
                    print("等待1秒后重试...")
                    time.sleep(1)
    else:
        print("没有发现临时录制文件需要清理")
    
    # 清理临时文件副本
    temp_copy = TEMP_MKV + '.copy'
    if os.path.exists(temp_copy):
        for i in range(3):
            try:
                os.unlink(temp_copy)
                print(f"已删除临时文件副本: {temp_copy}")
                break
            except Exception as e:
                print(f"清理临时文件副本时出错: {e}")
                if i < 2:  # 不是最后一次尝试
                    print("等待1秒后重试...")
                    time.sleep(1)
    else:
        print("没有发现临时文件副本需要清理")
        
    # 暂时注释掉清理流目录的代码（根据用户要求，保留生成的流文件）
    # try:
    #     if os.path.exists(STREAM_DIR):
    #         for file in os.listdir(STREAM_DIR):
    #             file_path = os.path.join(STREAM_DIR, file)
    #             if os.path.isfile(file_path):
    #                 os.unlink(file_path)
    #         print(f"已清理流目录: {STREAM_DIR}")
    # except Exception as e:
    #     print(f"清理流目录时出错: {e}")
    print("已按照用户要求，暂时保留流目录中的文件")

if __name__ == "__main__":
    # 检查必要的命令是否可用
    if not check_adb_available():
        print("错误: 未找到ADB工具。请确保Android SDK已安装并将adb添加到系统路径中。")
        sys.exit(1)
    
    if not check_command_available('scrcpy'):
        print("错误: 未找到scrcpy工具。请安装scrcpy并添加到系统路径中。")
        sys.exit(1)
    
    if not check_command_available('ffmpeg'):
        print("错误: 未找到ffmpeg工具。请安装ffmpeg并添加到系统路径中。")
        sys.exit(1)
    
    # 检查设备是否连接
    try:
        devices_result = subprocess.run(
            ['adb', 'devices'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        if 'device' not in devices_result.stdout:
            print("错误: 未检测到已连接的Android设备。请确保设备已通过USB连接并已启用USB调试。")
            sys.exit(1)
    except Exception as e:
        print(f"检查设备连接时出错: {e}")
        sys.exit(1)
    
    # 清理旧的临时文件
    cleanup()
    
    # 创建流目录
    create_stream_dir()
    
    # 启动各个进程
    processes = {}
    try:
        # 启动scrcpy和ffmpeg（通过管道直接传输数据）
        processes['scrcpy'], processes['ffmpeg'] = start_scrcpy_and_ffmpeg()
        
        # 检查scrcpy和ffmpeg是否成功启动
        if processes['scrcpy'] is None or processes['ffmpeg'] is None:
            print("错误: scrcpy或ffmpeg启动失败，无法继续执行")
            # 清理资源
            for process in processes.values():
                if process:
                    try:
                        process.terminate()
                    except Exception:
                        pass
            cleanup()
            sys.exit(1)
        
        time.sleep(2)  # 给ffmpeg一点时间启动
        
        # 启动WebSocket服务器
        processes['websocket'] = start_websocket_server()
        time.sleep(1)  # 给WebSocket服务器一点时间启动
        
        # 启动HTTP服务器
        processes['http'] = start_http_server()
        time.sleep(1)  # 给HTTP服务器一点时间启动
        
        print("\n=== 投屏服务已全部启动 ===")
        print(f"请在浏览器中访问: http://localhost:8000")
        print("功能说明:")
        print("1. 观看实时投屏")
        print("2. 点击屏幕可以控制手机")
        print("3. 点击'截图'按钮保存当前画面")
        print("4. 点击'全屏显示'切换全屏模式")
        print("\n按 Ctrl+C 停止所有服务...\n")
        
        # 监控进程
        monitor_processes(processes)
        
    except KeyboardInterrupt:
        print("\n正在停止所有服务...")
    finally:
        # 停止所有进程
        for process in processes.values():
            try:
                process.terminate()
                process.wait(timeout=5)
            except Exception:
                try:
                    process.kill()
                except Exception:
                    pass
        
        # 清理
        cleanup()
        
        print("所有服务已停止")